import os
import datetime
import pandas as pd
import akshare as ak
from analysis.hk_stock_dividend_batch import hk_stock_dividend_batch
import time
import random


def calc_latest_dividend_yield():
    cache_path = "../data/hk_stock_dividend/latest_close_cache.csv"

    # 一定要先初始化cache_df
    if os.path.exists(cache_path):
        cache_df = pd.read_csv(cache_path, dtype={"stock_code": str})
    else:
        cache_df = pd.DataFrame(columns=["stock_code", "latest_close"])

    csv_path = "../data/hk_stock_dividend/all_hk_stock_dividend.csv"
    if os.path.exists(csv_path):
        print("检测到all_hk_stock_dividend.csv，直接加载...")
        all_dividend_df = pd.read_csv(csv_path, encoding="utf-8")
        all_dividend_df.columns = all_dividend_df.columns.str.strip()
        print("all_dividend_df columns:", all_dividend_df.columns)
        # 如果实际列名不是stock_code，做重命名
        if "code" in all_dividend_df.columns:
            all_dividend_df.rename(columns={"code": "stock_code"}, inplace=True)
    else:
        print("未检测到all_hk_stock_dividend.csv，开始抓取数据...")
        all_dividend_df = hk_stock_dividend_batch()
        all_dividend_df.columns = all_dividend_df.columns.str.strip()
        print("all_dividend_df columns:", all_dividend_df.columns)

    # 只保留2024年分红数据
    df_2024 = all_dividend_df[all_dividend_df["年度"] == '2024'].copy()
    print("df_2024 shape:", df_2024.shape)
    print(df_2024)

    if df_2024.empty:
        print("没有2024年分红数据，返回空结果。")
        return pd.DataFrame(columns=["stock_code", "年度", "年度合计派息", "latest_close", "dividend_yield"])

    # 2. 保证stock_code为字符串且补齐5位
    df_2024["stock_code"] = df_2024["stock_code"].astype(str).str.zfill(5)

    # 3. 找出未缓存的股票代码
    cached_codes = set(cache_df["stock_code"])
    to_fetch_codes = [code for code in df_2024["stock_code"] if code not in cached_codes]

    # 4. 抓取未缓存的股票...
    # 获取每个股票最新一周的收盘价
    end_date = datetime.datetime.now().strftime("%Y%m%d")
    start_date = (datetime.datetime.now() - datetime.timedelta(days=30)).strftime("%Y%m%d")

    # 只抓取未缓存的
    new_close_list = []
    for code in to_fetch_codes:
        try:
            stock_weekly = ak.stock_hk_hist(
                symbol=code,
                adjust='qfq',
                period="weekly",
                start_date=start_date,
                end_date=end_date
            )
            latest_row = stock_weekly.iloc[-1]
            latest_close = latest_row["收盘"]
        except Exception as e:
            print(f"{code} 获取行情失败: {e}")
            latest_close = None
        new_close_list.append({"stock_code": code, "latest_close": latest_close})

        # 随机等待5-10秒
        sleep_time = random.uniform(1, 3)
        print(f"等待 {sleep_time:.2f} 秒后继续下一个请求...")
        time.sleep(sleep_time)

    # 合并新抓取和缓存，写回文件
    if new_close_list:
        new_close_df = pd.DataFrame(new_close_list)
        cache_df = pd.concat([cache_df, new_close_df], ignore_index=True)
        cache_df = cache_df.drop_duplicates(subset=["stock_code"], keep="last")
        cache_df.to_csv(cache_path, index=False, encoding="utf-8-sig")

    latest_close_df = cache_df

    # 合并分红和收盘价，计算股息率
    result = pd.merge(df_2024, latest_close_df, on="stock_code", how="left")
    result["dividend_yield"] = result["年度合计派息"] / result["latest_close"]
    result = result[["stock_code", "年度", "年度合计派息", "latest_close", "dividend_yield"]]

    return result


if __name__ == "__main__":
    df = calc_latest_dividend_yield()
    print(df)
    df.to_csv("../data/hk_stock_dividend/latest_dividend_yield.csv", index=False, encoding="utf-8-sig")
